package mr

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"regexp"
	"sort"
	"strconv"
	"time"
)
import "log"
import "net/rpc"
import "hash/fnv"

//
// Map functions return a slice of KeyValue.
//
type KeyValue struct {
	Key   string
	Value string
}

// for sorting by key.
type ByKey []KeyValue

// for sorting by key.
func (a ByKey) Len() int           { return len(a) }
func (a ByKey) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
func (a ByKey) Less(i, j int) bool { return a[i].Key < a[j].Key }

// 使用 ihash(key) % NReduce 为 Map 发出的每个 KeyValue 选择 reducer 任务编号。
func ihash(key string) int {
	h := fnv.New32a()
	h.Write([]byte(key))
	return int(h.Sum32() & 0x7fffffff)
}

func getReduceFile(id int) []string {
	ret := make([]string, 0, 0)
	reg := regexp.MustCompile(`mr-.*-` + strconv.Itoa(id))

	path, _ := os.Getwd()
	files, _ := ioutil.ReadDir(path)
	for _, file := range files {
		if reg.MatchString(file.Name()) {
			ret = append(ret, file.Name())
		}
	}
	return ret
}

func Worker(mapf func(string, string) []KeyValue, reducef func(string, []string) string) {
	args := GetTaskArgs{}
	sleepTime := 3 * time.Second

	for true {
		task := getTask(args).Task

		fmt.Printf("Get task successfully, task is %+v\n", task)

		if task == nil {
			fmt.Printf("GetTask faild, the task is nil\n")
			time.Sleep(sleepTime)
			continue
		}

		switch task.Type {
		case Exit:
			fmt.Printf("A worker exit....")
			os.Exit(0)
		case Sleep:
			time.Sleep(sleepTime)
		case Mapper:
			err := doMapper(mapf, task)
			if err == nil {
				reportDone(task)
			}
		case Reducer:
			err := doReducer(reducef, task)
			if err == nil {
				reportDone(task)
			}
		}
	}
}

func doMapper(mapf func(string, string) []KeyValue, task *Task) error {
	var intermediate []KeyValue
	filename := task.InputFile
	file, err := os.Open(filename)
	if err != nil {
		log.Fatalf("cannot open %v", filename)
		return err
	}

	content, err := ioutil.ReadAll(file)
	if err != nil {
		log.Fatalf("cannot read %v", filename)
		return err
	}

	file.Close()

	intermediate = mapf(filename, string(content))

	all := make([][]KeyValue, task.NReduce)

	for _, kv := range intermediate {
		all[ihash(kv.Key)%task.NReduce] = append(all[ihash(kv.Key)%task.NReduce], kv)
	}

	dir, _ := os.Getwd()
	for i := 0; i < task.NReduce; i++ {
		tmpFile, err := ioutil.TempFile(dir, "mr-tmp-*")

		if err != nil {
			continue
		}

		enc := json.NewEncoder(tmpFile)
		for _, kv := range all[i] {
			enc.Encode(kv)
		}
		tmpFile.Close()

		os.Rename(tmpFile.Name(), "mr-"+strconv.Itoa(task.ID)+"-"+strconv.Itoa(i))
	}

	return nil
}

func doReducer(reducef func(string, []string) string, task *Task) error {
	var intermediate []KeyValue
	// 整合键值对，参考讲义
	for _, fileName := range getReduceFile(task.ID) {
		file, _ := os.Open(fileName)
		dec := json.NewDecoder(file)
		for {
			var kv KeyValue
			if err := dec.Decode(&kv); err != nil {
				break
			}
			intermediate = append(intermediate, kv)
		}
		file.Close()
	}

	// sort
	sort.Sort(ByKey(intermediate))

	dir, _ := os.Getwd()
	tmpFile, err := ioutil.TempFile(dir, "mr-tmp-*")
	if err != nil {
		log.Fatal("Failed to create temp file", err)
		return err
	}

	i := 0
	for i < len(intermediate) {
		j := i + 1
		for j < len(intermediate) && intermediate[j].Key == intermediate[i].Key {
			j++
		}
		var values []string
		for k := i; k < j; k++ {
			values = append(values, intermediate[k].Value)
		}
		output := reducef(intermediate[i].Key, values)
		fmt.Fprintf(tmpFile, "%v %v\n", intermediate[i].Key, output)
		i = j
	}
	tmpFile.Close()
	os.Rename(tmpFile.Name(), "mr-out-"+strconv.Itoa(task.ID))
	return nil
}

func reportDone(task *Task) {
	args := ReportDoneArgs{Task: task}
	reply := ReportDoneReply{}

	ok := call("Coordinator.ReportDone", &args, &reply)
	if ok {
		fmt.Printf("Worker report task done, arg is %+v\n", args)
	} else {
		fmt.Printf("call failed!\n")
	}
}

func getTask(args GetTaskArgs) GetTaskReply {
	reply := GetTaskReply{}
	ok := call("Coordinator.DistributeTask", &args, &reply)
	if !ok {
		fmt.Printf("call failed!\n")
	}
	return reply
}

func CallExample() {
	args := ExampleArgs{}

	args.X = 99

	reply := ExampleReply{}

	ok := call("Coordinator.Example", &args, &reply)
	if ok {
		// reply.Y should be 100.
		fmt.Printf("reply.Y %v\n", reply.Y)
	} else {
		fmt.Printf("call failed!\n")
	}
}

//
// send an RPC request to the coordinator, wait for the response.
// usually returns true.
// returns false if something goes wrong.
//
func call(rpcname string, args interface{}, reply interface{}) bool {
	// c, err := rpc.DialHTTP("tcp", "127.0.0.1"+":1234")
	sockname := coordinatorSock()
	c, err := rpc.DialHTTP("unix", sockname)
	if err != nil {
		log.Fatal("dialing:", err)
	}
	defer c.Close()

	err = c.Call(rpcname, args, reply)
	if err == nil {
		return true
	}

	fmt.Println(err)
	return false
}
